import random
import pickle
import numpy as np
import tensorflow as tf
from tqdm import tqdm


def dofinetune(params):
    from sequence_to_sequence import SequenceToSequence
    from data_utils import batch_flow_bucket as batch_flow
    from threadgenerator import ThreadGenerator

    x_data, y_data = pickle.load(open('iChat.pkl', 'rb'))
    ws = pickle.load(open('ws.pkl', 'rb'))

    n_epoch = 5000000
    batch_size = 128

    steps = int(len(x_data) / batch_size) + 1

    config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False)

    # save_path = './model/s2s_iChat.ckpt'


    tf.reset_default_graph()
    with tf.Graph().as_default():
        random.seed(0)
        np.random.seed(0)
        tf.set_random_seed(0)

        with tf.Session(config=config) as sess:
            model_finetune = SequenceToSequence(input_vocab_size=len(ws), target_vocab_size=len(ws),
                                                batch_size=batch_size,
                                                **params)
            init = tf.global_variables_initializer()
            sess.run(init)
            model_finetune.finetune(sess, r'./model/s2s_iChat.ckpt.meta', r'./model/')

            flow = ThreadGenerator(
                batch_flow([x_data, y_data], ws, batch_size, add_end=[False, True]),
                queue_maxsize=30
            )

            for epoch in range(1, n_epoch + 1):
                costs = []
                save_path = './model/s2s_iChat.ckpt'
                bar = tqdm(range(steps), total=steps, desc='epoch {},loss=0.000000'.format(epoch))
                for _ in bar:
                    x, xl, y, yl = next(flow)
                    x = np.flip(x, axis=1)
                    cost, lr = model_finetune.train(sess, x, xl, y, yl, return_lr=True)
                    costs.append(cost)
                    bar.set_description('epoch {} loss={:.6f} lr={:.6f}'.format(
                        epoch,
                        np.mean(costs),
                        lr
                    ))
                model_finetune.save(sess, epoch, save_path)


def main():
    import json
    dofinetune(json.load(open('params.json')))


if __name__ == '__main__':
    main()
